#!/usr/bin/env python3
"""
检查YOLO模型信息
"""

from ultralytics import YOLO
import os

def check_yolo_model():
    print("🔍 检查YOLO模型信息")
    print("=" * 40)
    
    # 检查模型文件
    if os.path.exists("best.pt"):
        print("✅ 模型文件存在: best.pt")
        file_size = os.path.getsize("best.pt") / (1024*1024)
        print(f"📁 文件大小: {file_size:.2f} MB")
    else:
        print("❌ 模型文件不存在: best.pt")
        return
    
    try:
        # 加载模型
        print("🔄 正在加载模型...")
        model = YOLO("best.pt")
        print("✅ 模型加载成功")
        
        # 获取模型信息
        print("\n📊 模型信息:")
        print(f"模型类型: {type(model)}")
        
        # 获取类别信息
        if hasattr(model, 'names') and model.names:
            print(f"\n🏷️ 检测类别 ({len(model.names)} 个):")
            for i, name in model.names.items():
                print(f"  {i}: {name}")
        else:
            print("⚠️ 无法获取类别信息")
        
        # 测试模型预测
        print("\n🧪 测试模型预测...")
        import numpy as np
        
        # 创建一个测试图像（黑色背景）
        test_image = np.zeros((320, 320, 3), dtype=np.uint8)
        
        # 添加一些测试内容（白色矩形）
        test_image[100:200, 100:200] = [255, 255, 255]
        
        print("📸 使用测试图像进行预测...")
        results = model.predict(test_image, conf=0.1, verbose=False)
        
        print(f"预测结果数量: {len(results)}")
        for i, r in enumerate(results):
            if r.boxes is not None:
                print(f"结果 {i}: 检测到 {len(r.boxes)} 个目标")
                for j, box in enumerate(r.boxes):
                    conf = float(box.conf[0])
                    cls_id = int(box.cls[0])
                    label = model.names[cls_id] if cls_id in model.names else f"未知类别{cls_id}"
                    print(f"  目标 {j}: {label} (置信度: {conf:.3f})")
            else:
                print(f"结果 {i}: 未检测到目标")
        
        print("\n✅ YOLO模型检查完成")
        
    except Exception as e:
        print(f"❌ 模型检查失败: {e}")
        import traceback
        traceback.print_exc()

if __name__ == "__main__":
    check_yolo_model() 